# -*- coding: utf-8 -*-
# @Author: lidongdong
# @time  : 18-12-20 上午9:40
# @file  : mnist_loader.py


from mnist import MNIST
import numpy as np
import cv2
import random


class MnistLoader():
    def __init__(self, data_path, batch_size=64, single_channel=False):
        mnist = MNIST(data_path)
        images, labels = mnist.load_training()
        images = map(lambda x: np.asarray(x), images)

        images = map(lambda x: np.reshape(x, (28, 28)), images)
        images = map(lambda x: cv2.resize(x.astype(np.uint8), (64, 64)), images)

        images = map(lambda x: x.astype(np.uint8), images)
        images = map(lambda x: (x - 127.5) / 127.5, images)
        images = map(lambda x: np.expand_dims(x, -1), images)
        if not single_channel:
            images = map(lambda x: np.tile(x, [1, 1, 3]), images)
        cv2.imwrite("/home/jack/Desktop/dd.png", images[0])
        print images[0].shape
        images = map(lambda x: np.expand_dims(x, 0), images)
        images = map(lambda x: x.astype(np.float32), images)
        self.images = images
        self.batch_size = batch_size
        self.image_size = len(self.images)

    def get_batch(self):
        random.shuffle(self.images)
        for i in range(0, self.image_size, self.batch_size):
            batch_list = self.images[i: min(i + self.batch_size, self.image_size)]
            batch = np.concatenate(batch_list, axis=0)
            yield batch, None, None


if __name__ == '__main__':
    mnist_loader = MnistLoader("/home/jack/hdd_download/", 64)
    gen = mnist_loader.get_batch()
    for g in gen:
        print np.max(g[0]), np.min(g[0])
